package hook_proxy

import (
    "context"
    "encoding/binary"
    "fmt"
    "github.com/panjf2000/gnet"
    "go.uber.org/zap"
    "io"
    "net"
)

// @Desc: 具体项目要实现这个接口
type (
    // hook事件监听
    HookEvent interface {
    	//  获取要连接的服务器地址
        TryGetServerIpAddr(customContext interface{}) (string, error)
        // 收到消息的处理
        OnMessage(buf []byte, customCtx interface{}, toServer bool) []byte
        // 创建默认的用户自定义数据结构
        CreateEmptyCustomContext() interface{}
        // 下一条消息长度
        GetNextMsgLen(buf []byte) uint32
    }
    // @Desc: 默认的hookevent，为了让子类不用实现所有的接口
    DefaultHookEvent struct { }
)

// 获取服务器地址
func (e *DefaultHookEvent) TryGetServerIpAddr(customContext interface{}) (string, error) {
    return "127.0.0.1:25000", nil
}
// 消息处理
func (e *DefaultHookEvent) OnMessage(buf []byte, customCtx interface{}, toServer bool) []byte {
    return buf
}
// 创建默认的自定义数据
func (e *DefaultHookEvent) CreateEmptyCustomContext() interface{} {
    return nil
}

//  @Desc: 下一条消息长度,这里默认使用了小端前4字节
func (e *DefaultHookEvent) GetNextMsgLen(buf []byte) uint32 {
    if len(buf) < 4 {
        return 0
    }
    return binary.LittleEndian.Uint32(buf[:4]) + 4
}

//  @Desc: 代理类
type HookProxy struct {
    // 监听的客户端端口
    listenClientPort int32
    // 为了不实现所有函数,所以要用这个
    *gnet.EventServer
    // 事件回调
    hookEventer HookEvent
    // 根节点的关闭context
    rootCtx context.Context
}

// 客户端连接上来
func (hp *HookProxy) OnOpened(c gnet.Conn) (out []byte, action gnet.Action) {
    zap.S().Infof("收到连接[%s],[%d]", c.RemoteAddr().String(), hp.listenClientPort)
    pc := &ProxyContext{
        clientConn: c,
        stopChan:   make(chan bool, 1),
    }
    // 创建默认的自定义数据
    pc.customCtx = hp.hookEventer.CreateEmptyCustomContext()
    c.SetContext(pc)
    return
}

// 接收服务器消息
func (hp *HookProxy)RecvServerMessage(cancelCtx context.Context, pc *ProxyContext) {
    go func() {
        buffer := make([]byte, 0)
        readbuff := make([]byte, 1024000)
        for {
            msgLen, err := pc.serverConn.Read(readbuff)
            if err != nil {
                if err == io.EOF || msgLen < 1 {
                    pc.stopChan <- true
                    break
                }
            }
            buffer = append(buffer, readbuff[:msgLen]...)
            for {
                nextMsgLen := hp.hookEventer.GetNextMsgLen(buffer)
                if nextMsgLen < 1 ||
                    nextMsgLen > uint32(len(buffer)) {
                    break
                }
                afterBuf := hp.hookEventer.OnMessage(buffer[:nextMsgLen], pc.customCtx, false)
                buffer = buffer[nextMsgLen:]
                if len(afterBuf) < 1 {
                    zap.S().Warnf("[消息] afterbuf长度为0")
                    continue
                }
                if err = pc.clientConn.AsyncWrite(afterBuf); err != nil {
                    zap.S().Warnf("[消息] 发送给客户端失败[%v]", err)
                    pc.stopChan <- true
                }
            }
        }
    }()
    defer func() {
        if pc.clientConn != nil {
            pc.clientConn.Close()
        }
        if pc.serverConn != nil {
            pc.serverConn.Close()
        }
    }()
    select {
    case <-cancelCtx.Done():
    case <-pc.stopChan:
    }
}

// 客户端断开
func (hp *HookProxy) OnClosed(c gnet.Conn, err error) (action gnet.Action) {
    zap.S().Infof("[网络] 客户端[%s]关闭", c.RemoteAddr().String())
    if ctx, ok := c.Context().(*ProxyContext); ok {
        ctx.stopChan <- true
    }
    return
}

// 收到客户端消息
func (hp *HookProxy) React(frame []byte, c gnet.Conn) (out []byte, action gnet.Action) {
    if ctx, ok := c.Context().(*ProxyContext); ok {
        afterFrame := hp.hookEventer.OnMessage(frame, ctx.customCtx, true)
        if len(afterFrame) < 1 {
            zap.S().Errorf("[网络] afterframe长度为0")
            ctx.stopChan <- true
            return
        }
        if ctx.serverConn == nil {
            // 首次发消息，获取服务器地址。
            ipaddr, err := hp.hookEventer.TryGetServerIpAddr(ctx.customCtx)
            if err != nil {
                if err.Error() == "Wait" {
                    return
                }
                zap.S().Warnf("[网络] [%s]获取服务器地址失败", c.RemoteAddr().String())
                c.Close()
                return
            }
            zap.S().Info("获取到服务器地址", ipaddr)
            conn, err := net.Dial("tcp", ipaddr)
            if err != nil {
                zap.S().Warnf("[网络] 有连接连上来[%s]但是连接gate[%s]失败", c.RemoteAddr().String(), ipaddr)
                c.Close()
                return
            }
            zap.S().Infof("[网络] 有连接连上来[%s],连接[%s]成功,本地[%s]",
                c.RemoteAddr().String(),
                ipaddr,
                conn.LocalAddr().String())
            ctx.serverConn = conn
            go hp.RecvServerMessage(hp.rootCtx, ctx)
        }
        if n, err := ctx.serverConn.Write(afterFrame); err != nil || n < 1 { //转发
            zap.S().Warnf("[网络] 发送服务器[%d]字节,err[%v]", n, err)
            ctx.stopChan <- true
        }

    } else {
        zap.S().Warnf("[网络] [%s]没有上下文,非法", c.RemoteAddr().String())
        c.Close()
    }
    return
}

//  @Desc: 创建一个代理并且激活
//  @param ctx 进程主动关闭的上下文
//  @param port 监听客户端的端口
func NewHookProxyAndServe(ctx context.Context, port int32, eventer HookEvent) {
    hp := &HookProxy{
        listenClientPort: port,
        hookEventer:      eventer,
        rootCtx:          ctx,
    }
    err := gnet.Serve(hp, fmt.Sprintf("tcp://:%d", hp.listenClientPort),
        gnet.WithMulticore(false),
        gnet.WithReadBufferCap(1024000),
        gnet.WithSocketRecvBuffer(1024*1024*1024*10))
    if err != nil {
        zap.S().Errorf("[网络] 监听失败[%v]", err)
        // TODO 支持重连
        panic(fmt.Sprintf("监听端口[%v]失败", err))
    }
}
